/*
 * direct_tile.h
 *
 *  Created on: May 14, 2013
 *      Author: drewlewis
 */

#ifndef DIRECT_TILE_H_
#define DIRECT_TILE_H_

#include "tiled_array.h"
#include <chemistry/qc/basis/integral.h>
#include "integrals/TilePack.h"
#include "Global.h"
#include "integrals/integrals.hpp"
#include "integrals/range.hpp"
#include "integrals/tensor.hpp"
#include <TiledArray/eigen.h>

namespace TA = TiledArray;

// Tile of a tensor that directly computes integral elements
class direct_tile_4c {
    private:
        const TA::Array<double, 4, direct_tile_4c > *owner_;
        std::vector<std::size_t> index_;
        sc::IntegralEnginePool<sc::TwoBodyInt> *IntegralPool_;
        std::vector< atom_range > atom_indices_;


    public:
        typedef double value_type;
        typedef TA::Tensor<double> eval_type;

        direct_tile_4c(TA::Array<double, 4, direct_tile_4c> *owner,
                    const std::vector<std::size_t>& index,
                    sc::IntegralEnginePool<sc::TwoBodyInt> *IntegralPool,
                    std::vector< atom_range > &atom_indices):
                        owner_(owner), index_(index),
                        IntegralPool_(IntegralPool),
                        atom_indices_(atom_indices)
        {}

        //don't use
        direct_tile_4c(){};

        inline operator TA::Tensor<double> () const {
            TA::Range range = owner_->trange().make_tile_range(index_);
            TA::Tensor<double> tile(range); // output tile
            std::fill(tile.begin(), tile.end(), 0);

            sc::Ref<sc::TwoBodyInt> eri_= IntegralPool_->intance();

            Basis basis = eri_->basis();

            // Get list of shells for mpqc tensor object.
            shell_list P = get_shell_list(atom_indices_[0], basis);
            shell_list Q = get_shell_list(atom_indices_[1], basis);
            shell_list R = get_shell_list(atom_indices_[2], basis);
            shell_list S = get_shell_list(atom_indices_[3], basis);

            //Use a mpqc tensor to pack the tile
            const std::size_t dim[] = { range.size()[0], range.size()[1],
                                        range.size()[2], range.size()[3]};

            Tensor4 tile_map(tile.data(), dim);
            mpqc::integrals::evaluate(eri_, P, Q, R, S, tile_map);
            return tile;
        }

        template <typename Archive>
        void serialize(Archive& ar) {
            assert(false);
        }
};
#if 0
class direct_G_2c {
private:
    TA::Array<double, 2, direct_G_2c> *owner_;
    std::vector<std::size_t> index_;
    sc::IntegralEngines *engine_list_;
    std::vector< atom_range > atom_indices_;
    ARRAY2 *D_;

public:
    typedef double value_type;
    typedef TA::Tensor<double> eval_type;

    direct_G_2c(){};

    direct_G_2c(TA::Array<double, 2, direct_G_2c> *owner,
                std::vector<std::size_t> &index,
                sc::IntegralEngines *engine_list,
                std::vector< atom_range > &atom_indices,
                ARRAY2 *D):
                    owner_(owner), index_(index), engine_list_(engine_list),
                    atom_indices_(atom_indices), D_(D)
    {}


    inline operator TA::Tensor<double> () const {
        // Get tile size and location and make tile
        TA::Range range = owner_->trange().make_tile_range(index_);
        TA::Tensor<double> tile(range);

        // Get integral information
        sc::Ref<sc::TwoBodyInt> eri_ = engine_list_->get_engine<sc::ERI>();
        Basis basis_ = eri_->basis();

        std::vector<int> P = get_shell_list( atom_indices_[0], basis_);
        std::vector<int> Q = get_shell_list( atom_indices_[1], basis_);

        std::cout << "P = " << P << std::endl;
        std::cout << "Q = " << Q << std::endl;

        for(int i = 0; i < tile.size(); ++i){
            tile[i] = 1.438;
        }

        return tile;
    }

    template <typename Archive>
    void serialize(Archive& ar) {
        assert(false);
    }


};
#endif

class direct_coulomb_tile_3c {

private:
   const TA::Array<double, 3, direct_coulomb_tile_3c> *owner_;
   std::vector<std::size_t> index_;
   sc::IntegralEnginePool<sc::TwoBodyThreeCenterInt> *pool_;
   std::vector< atom_range > atom_indices_;


public:
   typedef double value_type;
   typedef TA::Tensor<double> eval_type;

   direct_coulomb_tile_3c(){};
   direct_coulomb_tile_3c(TA::Array<double, 3, direct_coulomb_tile_3c> *owner,
                       const std::vector<std::size_t> &index,
                       sc::IntegralEnginePool<sc::TwoBodyThreeCenterInt> *pool,
                       const std::vector< atom_range > &atom_indices):
                           owner_(owner), index_(index),
                           pool_(pool), atom_indices_(atom_indices)
   {}

   inline operator TA::Tensor<double> () const {
       TA::Range range = owner_->trange().make_tile_range(index_);
       TA::Tensor<double> tile(range); // output tile
       sc::Ref<sc::TwoBodyThreeCenterInt> three_c = pool_->intance();
       Basis basis_obs = three_c->basis1();
       Basis basis_df = three_c->basis3();

       // Get list of shells for mpqc tensor object.
       std::vector<int> P = get_shell_list( atom_indices_[0], basis_obs);
       std::vector<int> Q = get_shell_list( atom_indices_[1], basis_obs);
       std::vector<int> R = get_shell_list( atom_indices_[2], basis_df);

       //Use a mpqc tensor to pack the tile
       const std::size_t dim[] = { range.size()[0], range.size()[1],
                                   range.size()[2]};

       Tensor3 tile_map(tile.data(), dim);
       mpqc::integrals::evaluate(three_c, P, Q, R, tile_map);
       return tile;
   }

   template <typename Archive>
   void serialize(Archive &ar){
       assert(false);
   }
};


inline void direct_coulomb_3c_fill(TA::Array<double, 3, direct_coulomb_tile_3c> &array,
                                sc::IntegralEnginePool<sc::TwoBodyThreeCenterInt> &pool,
                                const cluster_atoms &clusters){
    TA::Array<double,2>::pmap_interface::const_iterator it =
                    array.get_pmap()->begin();
    TA::Array<double,2>::pmap_interface::const_iterator end =
                    array.get_pmap()->end();

    for(; it != end ; ++it){
        std::vector<std::size_t> index = array.range().idx(*it);

        std::vector< atom_range > atom_indices;
        for(std::size_t i = 0; i < index.size(); ++i)
            atom_indices.push_back( clusters[ index[i] ] );

        array.set(*it, direct_coulomb_tile_3c(&array, index, &pool, atom_indices));
    }
}

inline void four_center_int_fill(TA::Array<double, 4, direct_tile_4c>& array,
                          sc::IntegralEnginePool<sc::TwoBodyInt> &pool,
                          cluster_atoms &clusters){
    TA::Array<double, 4>::pmap_interface::const_iterator begin =
                    array.get_pmap()->begin();
    TA::Array<double, 4>::pmap_interface::const_iterator end =
                    array.get_pmap()->end();

    for(; begin != end; ++begin){
        //Get atoms on each tile
        std::vector<std::size_t> index = array.range().idx(*begin);

        std::vector< atom_range > atom_indices;
        for(std::size_t i = 0; i < index.size(); ++i)
            atom_indices.push_back( clusters[ index[i] ] );

        array.set(*begin, direct_tile_4c(&array, index, &pool, atom_indices) );
        array.get_world().gop.fence();
    }
}


#if 0
inline void direct_G_fill(TA::Array<double, 2, direct_G_2c> &array,
                          sc::IntegralEngines &engine_list,
                          cluster_atoms &clusters,
                          ARRAY2 &D){
    TA::Array<double, 2>::pmap_interface::const_iterator begin =
                    array.get_pmap()->begin();
    TA::Array<double, 2>::pmap_interface::const_iterator end =
                    array.get_pmap()->end();

    sc::Ref<sc::TwoBodyInt> eri_ = engine_list.get_engine<sc::ERI>();
    Matrix screen = two_e_abab(eri_);
    std::cout << "Screen = \n" << screen << std::endl;
    for(; begin != end; ++begin){
        //Get atoms on each tile
        std::vector<std::size_t> index = array.range().idx(*begin);

        std::vector< atom_range > atom_indices;
        for(std::size_t i = 0; i < index.size(); ++i)
            atom_indices.push_back( clusters[ index[i] ] );

        array.set(*begin, direct_G_2c(&array, index, &engine_list, atom_indices, &D) );
        array.get_world().gop.fence();
    }
}
#endif



#endif /* DIRECT_TILE_H_ */
